"""交易预测 Agent."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Dict, Optional

from ..prompts import build_prediction_prompt, load_task_template
from .base import AgentConfig, BaseETFAnalysisAgent


@dataclass
class PredictionContext:
    fund: str
    current_time: str
    position: str
    latest_quote: str
    quant_summary: str
    risk_summary: str
    technical_snapshot: str
    model_name: str
    model_prediction: str | None = None  # 模型预测结果


class PredictionAgent(BaseETFAnalysisAgent):
    """结合量化与风控结果生成交易预测报告."""

    def __init__(self, *, config: Optional[AgentConfig] = None) -> None:
        super().__init__(config=config)
        self._prompt = build_prediction_prompt()
        self._task_template = load_task_template("prediction_task")

    def build_inputs(self, context: PredictionContext) -> Dict[str, str]:
        inputs = {
            "current_time": context.current_time,
            "fund": context.fund,
            "position": context.position,
            "latest_quote": context.latest_quote,
            "quant_summary": context.quant_summary,
            "risk_summary": context.risk_summary,
            "technical_snapshot": context.technical_snapshot,
            "task_description": self._task_template.description,
        }
        # 如果有模型预测结果，添加到输入中
        if context.model_prediction:
            inputs["model_prediction"] = context.model_prediction
        return inputs

    def run(self, context: PredictionContext) -> str:
        inputs = self.build_inputs(context)
        inputs["task_description"] += f"\n当前使用的大模型名称：{context.model_name}"
        formatted = self._prompt.format(**inputs)
        return self.llm.invoke(formatted).content  # type: ignore[return-value]

